from paddleocr import PPStructure, save_structure_res
from pdf2image import convert_from_path
import os
import json
from datetime import datetime


class PDFOCRProcessor:
    def __init__(self, use_gpu=True, debug=True):
        """
        初始化OCR处理器
        :param use_gpu: 是否使用GPU
        :param debug: 是否打印调试信息
        """
        self.engine = PPStructure(show_log=True, use_gpu=use_gpu, layout=True, lang='ch')
        self.debug = debug

    def convert_pdf_to_images(self, pdf_path, output_dir=None):
        """
        将PDF转换为图像
        """
        if output_dir is None:
            output_dir = f"./paddleocr_temp_images_{datetime.now().strftime('%Y%m%d_%H%M%S')}"

        os.makedirs(output_dir, exist_ok=True)

        images = convert_from_path(pdf_path, dpi=600)
        image_paths = []

        for i, image in enumerate(images):
            image_path = os.path.join(output_dir, f'page_{i + 1}.jpg')
            image.save(image_path, 'JPEG')
            image_paths.append(image_path)

        return image_paths, output_dir

    def process_image(self, image_path):
        """
        处理单个图像
        """
        if self.debug:
            print(f"\n处理图像: {image_path}")

        result = self.engine(image_path)

        if self.debug:
            print(f"原始结果: {result[:2]}...")  # 只打印前两个结果避免过多输出
        for page in result:
            for item in page['res']:
                print(f"识别文字: {item['text']} (置信度: {item['confidence']})")
                print(f"边界框: {item['text_region']}")

        return result

    def save_results(self, results, output_path):
        """
        保存OCR结果
        """
        processed_results = []

        for page_num, page_result in enumerate(results, 1):
            if self.debug:
                print(f"\n处理第 {page_num} 页结果")
                print(f"页面结果数量: {len(page_result)}")

            page_data = {
                "page_number": page_num,
                "regions": []
            }

            for region in page_result:
                if self.debug:
                    print(f"区域类型: {region.get('type', 'unknown')}")

                # 根据区域类型分别处理
                if region.get('type') == 'title':
                    region_data = {
                        "type": "title",
                        "content": region.get('text', ''),
                        "bbox": region.get('bbox', [])
                    }
                elif region.get('type') == 'text':

                    for region_item in region['res']:
                        bbox_list = region_item.get('text_region', [[]])
                        region_data = {
                            "type": "text",
                            "content": region_item.get('text', ''),
                            "bbox": [item for sublist in bbox_list for item in sublist],
                            "confidence": region_item.get('confidence', 0)
                        }
                        page_data["regions"].append(region_data)
                elif region.get('type') == 'table':
                    # 处理表格内容
                    table_content = []
                    if 'res' in region and 'cell_contents' in region['res']:
                        for row in region['res']['cell_contents']:
                            processed_row = []
                            for cell in row:
                                if isinstance(cell, dict) and 'text' in cell:
                                    processed_row.append(cell['text'])
                                else:
                                    processed_row.append(str(cell))
                            table_content.append(processed_row)

                    region_data = {
                        "type": "table",
                        "content": table_content,
                        "bbox": region.get('bbox', [])
                    }
                else:
                    # 处理其他类型的内容
                    region_data = {
                        "type": region.get('type', 'unknown'),
                        "content": str(region.get('text', '')),
                        "bbox": region.get('bbox', [])
                    }

            processed_results.append(page_data)

        # 保存JSON结果
        with open(output_path + '.json', 'w', encoding='utf-8') as f:
            json.dump(processed_results, f, ensure_ascii=False, indent=2)

        # 保存文本结果
        # with open(output_path + '.txt', 'w', encoding='utf-8') as f:
        #     for page in processed_results:
        #         f.write(f"\n=== 第 {page['page_number']} 页 ===\n\n")
        #         for region in page["regions"]:
        #             if region["type"] in ["text", "title"]:
        #                 f.write(f"[{region['type']}] {region['content']}\n")
        #             elif region["type"] == "table":
        #                 f.write("\n【表格内容】:\n")
        #                 for row in region["content"]:
        #                     f.write("\t".join(row) + "\n")
        #             f.write("\n")

        if self.debug:
            print(f"\n处理了 {len(processed_results)} 页")
            print(f"结果已保存到: {output_path}.json 和 {output_path}.txt")

    def process_pdf(self, pdf_path, output_path):
        """
        处理PDF文件的主函数
        """
        try:
            if self.debug:
                print(f"\n开始处理PDF: {pdf_path}")

            os.makedirs(os.path.dirname(output_path), exist_ok=True)

            image_paths, temp_dir = self.convert_pdf_to_images(pdf_path)

            if self.debug:
                print(f"转换得到 {len(image_paths)} 个图像文件")

            all_results = []
            for image_path in image_paths:
                result = self.process_image(image_path)
                all_results.append(result)

            self.save_results(all_results, output_path)

            # 清理临时文件
            for image_path in image_paths:
                os.remove(image_path)
            os.rmdir(temp_dir)

            return output_path + '.json', output_path + '.txt'

        except Exception as e:
            print(f"处理PDF时发生错误: {str(e)}")
            raise


def main():
    # 使用示例
    processor = PDFOCRProcessor(use_gpu=False, debug=True)  # 启用调试输出
    pdf_path = "./input.pdf"  # 替换为您的PDF文件路径
    output_path = "./output/paddleocr/result"

    try:
        json_path, txt_path = processor.process_pdf(pdf_path, output_path)
        print(f"处理完成！")
        print(f"JSON结果保存在: {json_path}")
        print(f"文本结果保存在: {txt_path}")
    except Exception as e:
        print(f"处理失败: {str(e)}")


if __name__ == "__main__":
    main()
